from setuptools import setup
from torch.utils.cpp_extension import BuildExtension, CUDAExtension
import os

nvcc_flags = [
    '-O3', '-std=c++17',
    # '-U__CUDA_NO_HALF_OPERATORS__', '-U__CUDA_NO_HALF_CONVERSIONS__', '-U__CUDA_NO_HALF2_OPERATORS__',
    '-use_fast_math',
]

c_flags = ['-O3', '-std=c++17']

_src_path = os.path.dirname(os.path.abspath(__file__))

setup(
    name="mesh_coverage_loss",
    version="0.1",
    # packages=find_packages(), # slow
    ext_modules=[
        CUDAExtension(
            name="mesh_coverage_loss._C",
            sources=[os.path.join(_src_path, 'src', f) for f in [
                'cuda_ops.cu',
                'bindings.cpp',
                'forward.cu',
                'backward.cu', 
            ]],
            extra_compile_args={
                # 'cxx': c_flags,
                'nvcc': nvcc_flags,
            }
        ),
    ],
    cmdclass={
        'build_ext': BuildExtension
    },
    install_requires=[
        "torch>=1.0.0",
    ],
)